#include <string>
#include <dirent.h>
#include "SBB.h"
#include "utils.h"
#include <stdio.h>
#include <stdlib.h>
#include <utility>
#include <algorithm>
#include <iterator>
#include <cassert>
#include <pthread.h>
#include <unistd.h>
#include <semaphore.h>

using namespace std;

int TaskPipe[2];
int ResultPipe[2];
sem_t sem;
typedef pair<SBB const *, SBB const *> TASK;
void *worker(void *args);
vector<pthread_t> tids;

int SetupThreadPool(int nthreads) {
  assert(nthreads > 0);
  if (pipe(TaskPipe))
	goto out2;
  if (pipe(ResultPipe))
	goto out1;
  if (sem_init(&sem, 0, 0))
	goto out0;

  tids.reserve(nthreads);
  int i;
  for (i = 0; i < nthreads; ++i) {
	tids.push_back(pthread_t());
	if (pthread_create(&tids[i], NULL, worker, NULL)) {
	  tids.pop_back();
	  break;
	}
  }
  return i;

out0:
  close(ResultPipe[0]);
  close(ResultPipe[1]);
out1:
  close(TaskPipe[0]);
  close(TaskPipe[1]);
out2:
  return -1;
}

void TearDownThreadPool() {
  close(TaskPipe[1]);
  for (pthread_t const &tid : tids) {
	pthread_join(tid, NULL);
  }
  close(ResultPipe[0]);
  close(ResultPipe[1]);
  close(TaskPipe[0]);
  sem_destroy(&sem);
}

RACEPAIR CollectResults() {
  TASK poison(NULL, NULL);
  for (int i = 0; i < tids.size(); ++i)
	write(TaskPipe[1], &poison, sizeof(TASK));
  RACEPAIR totalRaces(0, 0);
  for (int i = 0; i < tids.size(); ++i) {
	RACEPAIR p;
	int ret = read(ResultPipe[0], &p, sizeof(p));
	if (ret != sizeof(p)) {
	  perror("pipe read: ");
	  break;
	}
	totalRaces.first += p.first;
	totalRaces.second += p.second;
  }
  for (int i = 0; i < tids.size(); ++i)
	sem_post(&sem);
  return totalRaces;
}

typedef pair<int, int> INTP;
bool CheckPS(const string &PSDir, bool parallel) {
  assert(parallel >= 0);
  cout << "Checking PS " << PSDir << '\n';
  UMAP nameToSBBs;
  DIR *dirp = opendir(PSDir.c_str());
  if (!dirp) return false;
  struct dirent *de;
  string path;
  path.reserve(PSDir.size() + sizeof("/THREADXX"));
  path = PSDir;
  path.append(PSDir.back() == '/' ? "THREADXX" : "/THREADXX");
  cout << "Reading trace data....";
  cout.flush();
  int nfiles = 0;
  while (de = readdir(dirp)) {
	if (de->d_type != DT_REG) continue;
	nfiles++;
	auto b =  path.rfind("THREAD");
	assert(b != string::npos);
	path.replace(begin(path) + b, end(path), de->d_name);
	ifstream infile(path);
	vector<SBB> *v = new vector<SBB>();
	int count = 0;
	while (1) {
	  SBB sbb = SBB(infile);
	  if (sbb.IsLast()) break;
	  sbb.Coalesce();
	  v->push_back(sbb);
#ifdef DEBUG
	  cout << count++ << (sbb.IsRecv() ? " RECV\n" : " POST\n");
	  cout << sbb;
	  cout << "------------------------" << endl;;
#endif
	}
	string foo = path + " 0"; // To make ThreadCode happy.
	string tid = std::move(ThreadCode(foo.c_str(), foo.c_str() + foo.size()).first);
	nameToSBBs[tid] = v;
  }
  cout << "Read " << nfiles << '\n';
  RACEPAIR totalRaces(0, 0);
  for (auto leftIt = begin(nameToSBBs); leftIt != end(nameToSBBs); ++leftIt) {
	for (auto rightIt = next(leftIt); rightIt != end(nameToSBBs); ++rightIt) {
	  auto thread1 = *leftIt, thread2 = *rightIt;
	  if (thread1.first == thread2.first) continue;
	  auto &leftV = *thread1.second;
	  auto &rightV = *thread2.second;

	  if (!parallel)
		printf("===RACE REPORT==(LEFT: %s[%d] <<==>> RIGHT: %s[%d])\n",
			thread1.first.c_str(), leftV.size(),
			thread2.first.c_str(), rightV.size());
	  for (int i = 0; i < leftV.size(); ++i) {
		for (int j = 0; j < rightV.size(); ++j) {
		  if (!parallel) {
			printf(COLOR_BLUE "LEFT SBB %d: %d <<==>> RIGHT SBB %d: %d\n" COLOR_RESET,
				i, leftV[i].size(), j, rightV[j].size());

			RACEPAIR sumary = DRCheck(leftV[i], rightV[j]);
			if (sumary.first + sumary.second) 
			  printf(COLOR_RED "Sumary: %d W-W races %d R-W/R-W races total: %d\n" COLOR_RESET,
				  sumary.first, sumary.second, sumary.first + sumary.second);
			else
			  printf("Sumary: %d races\n", 0);
			totalRaces.first += sumary.first;
			totalRaces.second += sumary.second;
		  } else {
			TASK t(&leftV[i], &rightV[j]);
			write(TaskPipe[1], &t, sizeof(TASK));
		  }
		}
	  }
	}
  }
  // Collect results.
  if (parallel) {
	totalRaces = CollectResults();
  }
  printf("######################Races: %d WW, %d W-R/R-W ######################\n",
	  totalRaces.first, totalRaces.second);

  // Free vectors in nameToSBBs
  for (auto const &p : nameToSBBs) {
	delete p.second;
  }
}

void *worker(void *args) {
  RACEPAIR p;
  for (;;) {
	TASK t;
	int ret = read(TaskPipe[0], &t, sizeof(TASK));
	if (!ret)
	  pthread_exit(NULL);
	if (t.first == NULL && t.second == NULL) {
	  write(ResultPipe[1], &p, sizeof(RACEPAIR));
	  p = RACEPAIR(0, 0);
	  sem_wait(&sem);
	  continue;
	}
	RACEPAIR r = DRCheck(*t.first, *t.second);
	p.first += r.first;
	p.second += r.second;
  }
}

int main(int argc, char **argv)
{
  if (argc < 2) {
	puts("Usage: ./a.out root-trace-directory/ [num of threads]");
	return 0;
  }
  int parallel = 0;
  if (argc == 3) {
	parallel = atoi(argv[2]);
	if (parallel < 1) {
	  puts("num of threads must > 1");
	  return 1;
	}
	if ((parallel = SetupThreadPool(parallel)) < 1) {
	  cout << "setup thread pool failed" << endl;
	  return 1;
	}
  }
  string rootDir = argv[1];
  if (*(end(rootDir) - 1) != '/')
	rootDir.push_back('/');
  DIR *dirp = opendir(rootDir.c_str());
  if (!dirp) {
	perror("opendir:");
	return 1;
  }
  struct dirent *de;
  while (de = readdir(dirp)) {
	string path = de->d_name;
	if (de->d_type != DT_DIR || path.find("PS")) continue;
	CheckPS(rootDir + path, parallel);
  }
  if (parallel)
	TearDownThreadPool();
  return 0;
}
